// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <utility>

#include <pollux/common/utils.h>
#include <pollux/exec/split.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_DECLARE_FLAG(int32_t, steps);

TURBO_DECLARE_FLAG(int32_t, duration_sec);

TURBO_DECLARE_FLAG(int32_t, batch_size);

TURBO_DECLARE_FLAG(int32_t, num_batches);

TURBO_DECLARE_FLAG(double, null_ratio);

TURBO_DECLARE_FLAG(bool, enable_spill);

TURBO_DECLARE_FLAG(int32_t, max_spill_level);

TURBO_DECLARE_FLAG(bool, enable_oom_injection);

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec {

class RowNumberFuzzerBase {
 public:
  explicit RowNumberFuzzerBase(
      size_t initialSeed,
      std::unique_ptr<ReferenceQueryRunner>);

  void run();

  virtual ~RowNumberFuzzerBase() = default;

 protected:
  bool isTableScanSupported(const TypePtr& type);

  // Runs one test iteration from query plans generations, executions and result
  // verifications.
  virtual void runSingleIteration() = 0;

  // Sets up the Dwrf reader/writer, serializers and Hive connector for the
  // fuzzers.
  void setupReadWrite();

  static VectorFuzzer::Options getFuzzerOptions() {
    VectorFuzzer::Options opts;
    opts.vectorSize = turbo::get_flag(FLAGS_batch_size);
    opts.stringVariableLength = true;
    opts.stringLength = 100;
    opts.nullRatio = turbo::get_flag(FLAGS_null_ratio);
    return opts;
  }

  void seed(size_t seed) {
    currentSeed_ = seed;
    vectorFuzzer_.reSeed(seed);
    rng_.seed(currentSeed_);
  }

  void reSeed() {
    seed(rng_());
  }

  int32_t randInt(int32_t min, int32_t max) {
    return fuzzer::rand(rng_, min, max);
  }

  // Validates the plan with input and result with the reference query runner.
  void validateExpectedResults(
      const core::PlanNodePtr& plan,
      const std::vector<RowVectorPtr>& input,
      const RowVectorPtr& result);

  struct PlanWithSplits {
    core::PlanNodePtr plan;
    std::vector<Split> splits;

    explicit PlanWithSplits(
        core::PlanNodePtr _plan,
        const std::vector<Split>& _splits = {})
        : plan(std::move(_plan)), splits(_splits) {}
  };

  // Executes a plan with spilling and oom injection possibly.
  RowVectorPtr execute(
      const PlanWithSplits& plan,
      const std::shared_ptr<memory::MemoryPool>& pool,
      bool injectSpill,
      bool injectOOM,
      const std::optional<std::string>& spillConfig = std::nullopt,
      int maxSpillLevel = -1);

  // Tests a plan by executing it with and without spilling. OOM injection
  // also might be done based on FLAG_enable_oom_injection.
  void testPlan(
      const PlanWithSplits& plan,
      int32_t testNumber,
      const RowVectorPtr& expected,
      const std::optional<std::string>& spillConfig);

  FuzzerGenerator rng_;
  size_t currentSeed_{0};

  std::shared_ptr<memory::MemoryPool> rootPool_{
      memory::memoryManager()->addRootPool(
          "rowNumberFuzzer",
          memory::kMaxMemory,
          memory::MemoryReclaimer::create())};
  std::shared_ptr<memory::MemoryPool> pool_{rootPool_->addLeafChild(
      "rowNumberFuzzerLeaf",
      true,
      memory::MemoryReclaimer::create())};
  std::shared_ptr<memory::MemoryPool> writerPool_{rootPool_->addAggregateChild(
      "rowNumberFuzzerWriter",
      memory::MemoryReclaimer::create())};
  VectorFuzzer vectorFuzzer_;
  std::unique_ptr<ReferenceQueryRunner> referenceQueryRunner_;
};

} // namespace kumo::pollux::exec
